import logging
import cv2
import os
import fastdeploy as fd
import traceback
from tqdm import tqdm


class FastSR:
    def __init__(self, model_path="", model_name="PP-MSVSR_reds_x4", runtime_option=None):

        self.model_file = os.path.join(model_path, model_name, "model.pdmodel")
        self.params_file = os.path.join(model_path, model_name, "model.pdiparams")
        if runtime_option is None:
            self.runtime_option = fd.RuntimeOption()
        else:
            # 注意传入的 runtime_option 必须是fd.RuntimeOption
            self.runtime_option = runtime_option

        if model_name == "PP-MSVSR_reds_x4":
            self.model = fd.vision.sr.PPMSVSR(
                model_file=self.model_file,
                params_file=self.params_file,
                runtime_option=self.runtime_option
            )
        elif model_name == "BasicVSR_reds_x4":
            self.model = fd.vision.sr.BasicVSR(
                model_file=self.model_file,
                params_file=self.params_file,
                runtime_option=self.runtime_option
            )
        else:
            self.model = fd.vision.sr.EDVR(
                model_file=self.model_file,
                params_file=self.params_file,
                runtime_option=self.runtime_option
            )

    def inference(self, imgs):
        try:
            results_batch = self.model.predict(imgs)
            return results_batch

        except Exception as e:
            with open("./error.log", "w") as f:
                f.write(traceback.format_exc())
                logging.critical("infrerence failed!")
                exit(0)

def SR_Video(model, output_path, progress_signal):
    logging.info('model SR start')
    # 设置tempfile 文件名字
    tempfile = 'temp.mp4'

    # 视频读取设置
    capture = cv2.VideoCapture(tempfile)
    video_fps = capture.get(cv2.CAP_PROP_FPS)  # 设置全局变量
    video_frame_count = capture.get(cv2.CAP_PROP_FRAME_COUNT)

    # 模型超分
    results = []
    imgs = []
    frame_count = 0

    # 因为每两帧应用一次模型，所以总帧数需要除以2
    total_frames = int(video_frame_count)

    for i in tqdm(range(total_frames)):
        _, frame = capture.read()
        if frame is not None:
            imgs.append(frame)
            frame_count += 1
        else:
            break

        if frame_count == 2:  # 每两帧应用一次模型
            results_batch = model.inference(imgs)

            results.extend(results_batch)
            imgs = []
            frame_count = 0

            # 更新进度条
            progress_signal.emit(str(int(100 * i / total_frames)))

    capture.release()

    logging.info('model SR finished')

    # 设置输出视频宽高
    out_width = 1280
    out_height = 720

    logging.info(f"fps: {video_fps}\tframe_count: {len(results)}")

    # Create VideoWriter for output
    output_file = os.path.join(output_path, 'output.mp4')
    fucc = cv2.VideoWriter_fourcc(*"mp4v")
    video_out = cv2.VideoWriter(output_file, fucc, video_fps,
                                (out_width, out_height), True)

    if not video_out.isOpened():
        logging.critical("create video writer failed!")

    # 视频写到磁盘
    for item in results:
        video_out.write(item)
    logging.info("write finished, output video saved at: " + output_path)
    video_out.release()
